package edge

import (
	"context"
	"fmt"

	"github.com/DataDog/KubeHound/pkg/kubehound/graph/adapter"
	"github.com/DataDog/KubeHound/pkg/kubehound/graph/types"
	"github.com/DataDog/KubeHound/pkg/kubehound/models/converter"
	"github.com/DataDog/KubeHound/pkg/kubehound/storage/cache"
	"github.com/DataDog/KubeHound/pkg/kubehound/storage/storedb"
	"github.com/DataDog/KubeHound/pkg/kubehound/store/collections"
	gremlin "github.com/apache/tinkerpop/gremlin-go/v3/driver"
	"go.mongodb.org/mongo-driver/bson"
)

const (
	RoleBindCrbCrRName = "RoleBindClusteRoleBindingbClusterRoleRole"
)

func init() {
	Register(&RoleBindCrbCrR{}, RegisterDefault)
}

type RoleBindCrbCrR struct {
	BaseEdge
}

func (e *RoleBindCrbCrR) Label() string {
	return RoleBindLabel
}

func (e *RoleBindCrbCrR) Name() string {
	return RoleBindCrbCrRName
}

func (e *RoleBindCrbCrR) AttckTechniqueID() AttckTechniqueID {
	return AttckTechniqueValidAccounts
}

func (e *RoleBindCrbCrR) AttckTacticID() AttckTacticID {
	return AttckTacticPrivilegeEscalation
}

func (e *RoleBindCrbCrR) Processor(ctx context.Context, oic *converter.ObjectIDConverter, entry any) (any, error) {
	typed, ok := entry.(*roleBindGroup)
	if !ok {
		return nil, fmt.Errorf("invalid type passed to processor: %T", entry)
	}

	psid, err := oic.GraphID(ctx, typed.PermissionSet.Hex())
	if err != nil {
		return nil, fmt.Errorf("%s edge PermissionSet id convert: %w", e.Label(), err)
	}

	return psid, nil
}

func (e *RoleBindCrbCrR) Traversal() types.EdgeTraversal {
	return func(source *gremlin.GraphTraversalSource, inserts []any) *gremlin.GraphTraversal {
		g := source.GetGraphTraversal()

		if e.cfg.LargeClusterOptimizations {
			// For larger clusters simply target specific roles to reduce number of attack paths
			g.V().
				Has("runID", e.runtime.RunID.String()).
				Has("cluster", e.runtime.Cluster.Name).
				Has("class", "PermissionSet").
				Has("isNamespaced", true).
				// Temporary measure, until we scan and flag for sensitive roles
				Has("critical", true).
				// Has("role", P.Within(sensitiveRoles)).
				As("r").
				V(inserts...).
				Has("critical", false).
				AddE(e.Label()).
				To("r").
				Property("attckTechniqueID", string(e.AttckTechniqueID())).
				Property("attckTacticID", string(e.AttckTacticID())).
				Barrier().Limit(0)
		} else {
			// In smaller clusters we can still show the (large set of) attack paths generated by this attack
			g.V().
				Has("runID", e.runtime.RunID.String()).
				Has("cluster", e.runtime.Cluster.Name).
				Has("class", "PermissionSet").
				Has("isNamespaced", true).
				As("i").
				V(inserts...).
				Has("critical", false).
				AddE(e.Label()).
				To("i").
				Property("attckTechniqueID", string(e.AttckTechniqueID())).
				Property("attckTacticID", string(e.AttckTacticID())).
				Barrier().Limit(0)
		}

		return g
	}
}

func (e *RoleBindCrbCrR) Stream(ctx context.Context, store storedb.Provider, c cache.CacheReader,
	callback types.ProcessEntryCallback, complete types.CompleteQueryCallback) error {

	permissionSets := adapter.MongoDB(ctx, store).Collection(collections.PermissionSetName)
	// Handle clusterrolebindings against clusterroles
	pipeline := []bson.M{
		// $match stage
		{
			"$match": bson.M{
				"is_namespaced":        false,
				"runtime.runID":        e.runtime.RunID.String(),
				"runtime.cluster.name": e.runtime.Cluster.Name,
				"$and": []bson.M{
					// Looking for RBAC objects
					{
						"rules": bson.M{
							"$elemMatch": bson.M{
								"$or": []bson.M{
									{"apigroups": "*"},
									{"apigroups": "rbac.authorization.k8s.io"},
								},
							},
						},
					},
					// Looking for creation of rolebindings
					{
						"rules": bson.M{
							"$elemMatch": bson.M{
								"$and": []bson.M{
									{
										"$or": []bson.M{
											{"verbs": "create"},
											{"verbs": "*"},
										},
									},
									{
										"$or": []bson.M{
											{"resources": "rolebindings"},
											{"resources": "*"},
										},
									},
									{
										"$or": []bson.M{
											{"resourcenames": nil},
										},
									},
								},
							},
						},
					},
					// Looking for binding roles
					{
						"rules": bson.M{
							"$elemMatch": bson.M{
								"$and": []bson.M{
									{
										"$or": []bson.M{
											{"verbs": "bind"},
											{"verbs": "*"},
										},
									},
									{
										"$or": []bson.M{
											{"resources": "roles"},
											{"resources": "*"},
										},
									},
									{
										"$or": []bson.M{
											{"resourcenames": nil},
										},
									},
								},
							},
						},
					},
				},
			},
		},
		// Looking for all permissionSets link to the same namespace
		{
			"$lookup": bson.M{
				"from":         "rolebindings",
				"localField":   "role_binding_id",
				"foreignField": "_id",
				"as":           "rb",
			},
		},
		{
			"$unwind": bson.M{
				"path": "$rb",
			},
		},
		{
			"$match": bson.M{
				"rb.is_namespaced":     false,
				"runtime.runID":        e.runtime.RunID.String(),
				"runtime.cluster.name": e.runtime.Cluster.Name,
			},
		},
		// $project stage
		{
			"$project": bson.M{
				"_id": 1,
			},
		},
	}
	cur, err := permissionSets.Aggregate(ctx, pipeline)
	if err != nil {
		return err
	}
	defer cur.Close(ctx)

	return adapter.MongoCursorHandler[roleBindGroup](ctx, cur, callback, complete)
}
